import matplotlib.pyplot as plt
import numpy as np
import paddle
from paddle.nn import functional as F
from paddle.vision.transforms import transforms as T
from PIL import Image as PilImage

from cgdataset import CgDataSet
from model import CgNet


def test(image):
    flag_c = False
    flag_t = False
    C_count = 0
    T_count = 0
    for i in range(720):
        for j in range(720):
            if image[i, j] == 3:
                C_count = C_count + 1
            if image[i, j] == 4:
                T_count = T_count + 1
    if C_count > 10:
        flag_c = True
    if T_count > 10:
        flag_t = True
    
    if flag_c and flag_t:
        return "positive"
    elif flag_c and not flag_t:
        return "negative"
    else:
        return "invalid"

def data_predict(filepath):
    num_classes = 5
    network = CgNet(num_classes)
    model = paddle.Model(network)

    paddle.device.set_device('gpu:0')

    model.load('checkpoint/final.pdparams')
    model.prepare()

    predict_dataset = CgDataSet(filepath=filepath, mode='predict')
    predict_results = model.predict(predict_dataset)

    data = predict_results[0][0][0].transpose((1, 2, 0))
    mask = np.argmax(data, axis=-1)

    result = test(mask)

    return result

